#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
PostgreSQL配置优化模块
提供基于硬件资源的PostgreSQL配置优化功能
"""

import os
import psutil
import math
from typing import Dict, Any, Optional, Tuple
from ...core.utils import logger, run_command, detect_postgresql_version

class PostgresqlOptimizer:
    """PostgreSQL配置优化类"""
    
    def __init__(self):
        self.pg_version = detect_postgresql_version()
        
    def get_system_resources(self) -> Dict[str, Any]:
        """
        获取系统硬件资源信息
        
        Returns:
            包含CPU和内存信息的字典
        """
        try:
            # 获取CPU信息
            cpu_count = psutil.cpu_count(logical=False)  # 物理核心数
            cpu_logical = psutil.cpu_count(logical=True)  # 逻辑核心数
            
            # 获取内存信息
            memory = psutil.virtual_memory()
            total_memory_gb = memory.total / (1024**3)  # 转换为GB
            available_memory_gb = memory.available / (1024**3)
            
            # 获取磁盘信息（根目录）
            disk = psutil.disk_usage('/')
            total_disk_gb = disk.total / (1024**3)
            free_disk_gb = disk.free / (1024**3)
            
            return {
                'cpu_physical_cores': cpu_count,
                'cpu_logical_cores': cpu_logical,
                'total_memory_gb': round(total_memory_gb, 2),
                'available_memory_gb': round(available_memory_gb, 2),
                'total_disk_gb': round(total_disk_gb, 2),
                'free_disk_gb': round(free_disk_gb, 2),
                'system_type': 'Windows' if os.name == 'nt' else 'Linux'
            }
        except Exception as e:
            logger.error(f"获取系统资源信息失败: {e}")
            return {}
    
    def calculate_optimal_config(self, resources: Dict[str, Any]) -> Dict[str, Any]:
        """
        根据系统资源计算最优PostgreSQL配置
        
        Args:
            resources: 系统资源信息
            
        Returns:
            优化后的配置参数字典
        """
        if not resources:
            return {}
            
        # 基础配置参数
        config = {}
        
        # 内存相关配置
        total_memory_gb = resources['total_memory_gb']
        cpu_cores = resources['cpu_physical_cores']
        
        # shared_buffers: 总内存的25%左右，但不超过8GB
        shared_buffers = min(int(total_memory_gb * 0.25 * 1024), 8192)
        config['shared_buffers'] = f"{shared_buffers}MB"
        
        # effective_cache_size: 总内存的50%左右
        effective_cache_size = int(total_memory_gb * 0.5 * 1024)
        config['effective_cache_size'] = f"{effective_cache_size}MB"
        
        # maintenance_work_mem: 维护操作内存，根据总内存调整
        if total_memory_gb <= 4:
            maintenance_work_mem = 256
        elif total_memory_gb <= 8:
            maintenance_work_mem = 512
        elif total_memory_gb <= 16:
            maintenance_work_mem = 1024
        else:
            maintenance_work_mem = 2048
        config['maintenance_work_mem'] = f"{maintenance_work_mem}MB"
        
        # work_mem: 每个查询的工作内存，根据内存大小调整
        if total_memory_gb <= 4:
            work_mem = 8
        elif total_memory_gb <= 8:
            work_mem = 16
        elif total_memory_gb <= 16:
            work_mem = 32
        else:
            work_mem = 64
        config['work_mem'] = f"{work_mem}MB"
        
        # CPU相关配置
        # max_connections: 根据内存和CPU核心数调整
        if total_memory_gb <= 2:
            max_connections = 50
        elif total_memory_gb <= 4:
            max_connections = 100
        elif total_memory_gb <= 8:
            max_connections = 200
        elif total_memory_gb <= 16:
            max_connections = 300
        else:
            max_connections = 400
        
        # 根据CPU核心数调整连接数
        cpu_based_connections = cpu_cores * 16
        max_connections = min(max_connections, cpu_based_connections)
        
        config['max_connections'] = str(max_connections)
        
        # 并行查询相关配置
        config['max_parallel_workers'] = str(cpu_cores)
        config['max_parallel_workers_per_gather'] = str(max(1, cpu_cores // 2))
        
        # WAL相关配置
        # wal_buffers: 总内存的3%左右，最小16MB
        wal_buffers = max(16, int(total_memory_gb * 0.03 * 1024))
        config['wal_buffers'] = f"{wal_buffers}MB"
        
        # checkpoint相关配置
        config['checkpoint_completion_target'] = '0.9'
        config['wal_compression'] = 'on'
        
        # 日志相关配置
        config['log_min_duration_statement'] = '2000'  # 记录超过2秒的查询
        config['log_checkpoints'] = 'on'
        
        # 查询优化相关配置
        config['random_page_cost'] = '1.1'  # SSD硬盘优化
        config['effective_io_concurrency'] = '200'  # SSD优化
        config['seq_page_cost'] = '1.0'
        
        # VACUUM相关配置
        config['autovacuum_vacuum_scale_factor'] = '0.01'
        config['autovacuum_analyze_scale_factor'] = '0.005'
        config['autovacuum_max_workers'] = str(max(3, cpu_cores // 2))
        
        # 统计信息收集
        config['track_activities'] = 'on'
        
        return config
    
    def get_postgresql_conf_path(self) -> Optional[str]:
        """
        获取PostgreSQL配置文件路径
        
        Returns:
            配置文件路径，如果找不到则返回None
        """
        if not self.pg_version:
            logger.error("未检测到PostgreSQL版本")
            return None
            
        # 常见的PostgreSQL配置文件路径
        common_paths = []
        
        if os.name == 'nt':  # Windows
            common_paths = [
                f"C:\\Program Files\\PostgreSQL\\{self.pg_version}\\data\\postgresql.conf",
                f"C:\\PostgreSQL\\{self.pg_version}\\data\\postgresql.conf",
                f"D:\\PostgreSQL\\{self.pg_version}\\data\\postgresql.conf"
            ]
        else:  # Linux
            common_paths = [
                f"/etc/postgresql/{self.pg_version}/main/postgresql.conf",
                f"/var/lib/pgsql/{self.pg_version}/data/postgresql.conf",
                f"/usr/local/pgsql/data/postgresql.conf",
                f"/opt/PostgreSQL/{self.pg_version}/data/postgresql.conf"
            ]
        
        # 检查路径是否存在
        for path in common_paths:
            if os.path.exists(path):
                return path
                
        # 尝试通过pg_config命令获取
        try:
            result = run_command(['pg_config', '--sysconfdir'], capture_output=True)
            if result.returncode == 0:
                config_dir = result.stdout.strip()
                conf_path = os.path.join(config_dir, 'postgresql.conf')
                if os.path.exists(conf_path):
                    return conf_path
        except:
            pass
            
        logger.error("无法找到PostgreSQL配置文件")
        return None
    
    def backup_config(self, config_path: str) -> bool:
        """
        备份原始配置文件
        
        Args:
            config_path: 配置文件路径
            
        Returns:
            备份是否成功
        """
        try:
            backup_path = f"{config_path}.backup"
            
            # 如果备份文件已存在，添加时间戳
            if os.path.exists(backup_path):
                import time
                timestamp = time.strftime("%Y%m%d_%H%M%S")
                backup_path = f"{config_path}.backup_{timestamp}"
            
            # 复制文件
            import shutil
            shutil.copy2(config_path, backup_path)
            logger.info(f"配置文件已备份到: {backup_path}")
            return True
        except Exception as e:
            logger.error(f"备份配置文件失败: {e}")
            return False
    
    def optimize_postgresql_config(self) -> bool:
        """
        一键优化PostgreSQL配置
        
        Returns:
            优化是否成功
        """
        try:
            # 获取系统资源信息
            resources = self.get_system_resources()
            if not resources:
                logger.error("无法获取系统资源信息")
                return False
                
            # 显示当前系统资源
            print("当前系统资源信息:")
            print(f"CPU 物理核心数: {resources['cpu_physical_cores']}")
            print(f"CPU 逻辑核心数: {resources['cpu_logical_cores']}")
            print(f"总内存: {resources['total_memory_gb']} GB")
            print(f"可用内存: {resources['available_memory_gb']} GB")
            print(f"系统类型: {resources['system_type']}")
            print()
            
            # 计算优化配置
            optimal_config = self.calculate_optimal_config(resources)
            if not optimal_config:
                logger.error("无法计算优化配置")
                return False
                
            # 获取配置文件路径
            config_path = self.get_postgresql_conf_path()
            if not config_path:
                logger.error("无法找到PostgreSQL配置文件")
                return False
                
            # 备份原始配置
            if not self.backup_config(config_path):
                logger.error("备份配置文件失败")
                return False
                
            # 读取现有配置
            with open(config_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # 构建新的配置内容
            new_lines = []
            modified_lines = []
            
            for line in lines:
                stripped_line = line.strip()
                
                # 跳过注释和空行
                if not stripped_line or stripped_line.startswith('#'):
                    new_lines.append(line)
                    continue
                    
                # 检查是否是配置项
                if '=' in stripped_line:
                    key = stripped_line.split('=')[0].strip()
                    
                    # 如果是需要优化的配置项
                    if key in optimal_config:
                        # 替换配置值
                        new_line = f"{key} = {optimal_config[key]}\n"
                        new_lines.append(new_line)
                        modified_lines.append((key, optimal_config[key]))
                        # 从待处理配置中移除
                        del optimal_config[key]
                    else:
                        new_lines.append(line)
                else:
                    new_lines.append(line)
            
            # 添加新的配置项
            if optimal_config:
                new_lines.append("\n# 自动优化的配置项\n")
                for key, value in optimal_config.items():
                    new_lines.append(f"{key} = {value}\n")
                    modified_lines.append((key, value))
            
            # 写入新配置
            with open(config_path, 'w', encoding='utf-8') as f:
                f.writelines(new_lines)
            
            # 显示修改的配置
            print("优化的配置参数:")
            for key, value in modified_lines:
                print(f"  {key} = {value}")
            
            print(f"\n配置文件已优化: {config_path}")
            print("请重启PostgreSQL服务以使配置生效")
            
            return True
            
        except Exception as e:
            logger.error(f"优化PostgreSQL配置失败: {e}")
            return False
    
    def show_current_config(self) -> bool:
        """
        显示当前PostgreSQL配置
        
        Returns:
            显示是否成功
        """
        try:
            config_path = self.get_postgresql_conf_path()
            if not config_path:
                logger.error("无法找到PostgreSQL配置文件")
                return False
                
            print(f"当前PostgreSQL配置文件路径: {config_path}")
            print("\n主要配置参数:")
            
            # 读取配置文件
            with open(config_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # 显示主要配置参数
            important_params = [
                'shared_buffers', 'effective_cache_size', 'maintenance_work_mem',
                'work_mem', 'max_connections', 'max_parallel_workers',
                'wal_buffers', 'checkpoint_completion_target'
            ]
            
            for line in lines:
                stripped_line = line.strip()
                if stripped_line and not stripped_line.startswith('#') and '=' in stripped_line:
                    key = stripped_line.split('=')[0].strip()
                    if key in important_params:
                        print(f"  {key} = {stripped_line.split('=')[1].strip()}")
            
            return True
            
        except Exception as e:
            logger.error(f"显示当前配置失败: {e}")
            return False
    
    def restore_config(self) -> bool:
        """
        恢复配置文件备份
        
        Returns:
            恢复是否成功
        """
        try:
            config_path = self.get_postgresql_conf_path()
            if not config_path:
                logger.error("无法找到PostgreSQL配置文件")
                return False
                
            backup_path = f"{config_path}.backup"
            if not os.path.exists(backup_path):
                logger.error("找不到配置文件备份")
                return False
                
            # 恢复备份
            import shutil
            shutil.copy2(backup_path, config_path)
            logger.info("配置文件已恢复到备份状态")
            return True
            
        except Exception as e:
            logger.error(f"恢复配置文件失败: {e}")
            return False